/*
 * Copyright (C) 2018 The Asiainfo-Ability Authors
 *
 *      Licensed under the Apache License, Version 2.0 (the "License");
 *      you may not use this file except in compliance with the License.
 *      You may obtain a copy of the License at
 *
 *          http://www.apache.org/licenses/LICENSE-2.0
 *
 *      Unless required by applicable law or agreed to in writing, software
 *      distributed under the License is distributed on an "AS IS" BASIS,
 *      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *      See the License for the specific language governing permissions and
 *      limitations under the License.
 */

package org.asiainfo.ability.provider.custom;

/**
 * @Author: visen
 * @Date: 2018/3/20
 * @Description:
 */

import org.springframework.security.web.firewall.FirewalledRequest;

import java.io.*;

import javax.servlet.ServletOutputStream;
import javax.servlet.WriteListener;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;


public class ResponseWrapper extends HttpServletResponseWrapper {


    /**
     * 输出字节流的对象，用于数据输出
     */
    private ByteArrayOutputStream outstream;

    /**
     * 输出字符串文本的对象，用与JSP页面的输出
     */
    private CharArrayWriter output;

    /**
     * 返回字符串格式数据
     *
     * @return String
     */
    public String toString() {
        return output.toString();
    }

    /**
     * 返回字节数组格式的数据
     *
     * @return byte[]
     */
    public byte[] toByteArray() {
        return outstream.toByteArray();
    }

    /**
     * @param response HttpServletResponse
     */
    private HttpServletResponse response;

    public ResponseWrapper(HttpServletResponse response) {
        super(response);

        outstream = new ByteArrayOutputStream(8192);
    }

    @Override
    public ServletOutputStream getOutputStream() throws IOException {
        return new WrapperOutputStream(outstream);
    }

    private PrintWriter printWriter;

    @Override
    public PrintWriter getWriter() throws IOException {
        if(printWriter == null) {
            printWriter = new PrintWriter(new OutputStreamWriter(outstream, super.getCharacterEncoding()));
        }

        return printWriter;
    }

    public void reset() {
        this.outstream.reset();
    }

    public void writeError(String error) {
        PrintWriter printWriter = null;
        try {
            printWriter = getWriter();
            printWriter.write(error);
            printWriter.flush();
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            if(outstream != null) {
                try {
                    outstream.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }
    }
    public void flushResponse() {

        ServletOutputStream output = null;
        try {
            output = this.getResponse().getOutputStream();
            outstream.writeTo(output);
            output.flush();
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            try {
                if (output != null)
                    output.close();
                if(printWriter != null) {
                    printWriter.close();
                }
            } catch (IOException e) {
                e.printStackTrace();
            }
        }


    }



    class WrapperOutputStream extends ServletOutputStream {
        private ByteArrayOutputStream bos;

        public WrapperOutputStream(ByteArrayOutputStream bos) {
            this.bos = bos;
        }

        @Override
        public void write(int b)
                throws IOException {
            bos.write(b);
        }


        @Override
        public boolean isReady() {

            return false;

        }

        @Override
        public void setWriteListener(WriteListener arg0) {

            // TODO Auto-generated method stub

        }
    }

}